import os
import time

import torch
import torchvision
from torch import nn
from torch.autograd import Variable
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import MNIST, CIFAR10
from torchvision.utils import save_image
import numpy as np
import pandas as pd
import torch
import os
from datetime import datetime
import matplotlib.pyplot as plt
from torch.utils.data import DataLoader
from torch.autograd import Variable
from torch import nn
from torch import optim
from torchvision import datasets
import sys
import math
from sklearn import datasets
from sklearn.model_selection import train_test_split
from sklearn.model_selection import GridSearchCV
#from sklearn import cross_validation,metrics
from sklearn import model_selection as cv
from sklearn import metrics
from sklearn.model_selection import KFold, cross_val_score





def getData():
    np.random.seed(10) # seed相同, 则每次获得的矩阵都是相同的
    endLoc = 1000  # 终止点, 一共1500个点不改变, 这个终止点只是为了限制点之间的间隔
    totalPoints = 1500
    xMat = []
    yMat = []
    for i in range(totalPoints):
        t = np.random.random(1)[0]*2 # 正态分布*2
        x1Temp = t**2 - t + 1
        x2Temp = math.sin(t)
        #x3Temp = t**3 + t
        x4Temp = 2*math.cos(0.08*t)*math.sin(0.06*t)
        x5Temp = math.sin(0.4*t) + 1.5*math.cos(0.2*t)
        #xMat.append([x1Temp, x2Temp, x3Temp, x4Temp, x5Temp])
        xMat.append([x1Temp, x2Temp, x4Temp, x5Temp])

        #miu = math.exp(0.7*x1Temp + 0.2*x2Temp + 0.4*x3Temp + 0.7*x4Temp + 0.5*x5Temp)
        miu = math.exp(0.7*x1Temp + 0.2*x2Temp + 0.7*x4Temp + 0.5*x5Temp)
        yToAdd = np.random.poisson(miu)
        yMat.append(yToAdd)

    xMat, yMat = np.array(xMat), np.array(yMat)
    print("过程变量形状为: ", xMat.shape, "质量变量形状为: ", yMat.shape)
    return xMat, yMat

def getNoisData():
    # 返回同样的xMat, yMat, 加了噪声
    xMat, yMat = getData()
    # print(yMat.dtype)
    xMat += np.random.rand(xMat.shape[0], xMat.shape[1])*0.1
    return xMat, yMat

def addNoise(x):
    # 注意x是带batchSize这个维度的
    return x + torch.randn(x.shape[0], x.shape[1])*0.1

def getDataLoader(historyRecordNum = 1000, isNoisy = False):
    batchSize = 128  # 128

    # 用来训练初始网络的历史数据数量
    torch.set_printoptions(precision=5)
    torch.manual_seed(233)

    if isNoisy:
        xsDf, ysDf = getNoisData()
    else:
        xsDf, ysDf = getData()

    xTrain = torch.tensor(xsDf[:historyRecordNum, :], dtype=torch.float32)
    xPredi = torch.tensor(xsDf[historyRecordNum:, :], dtype=torch.float32)
    yTrain = torch.tensor(ysDf[:historyRecordNum, ], dtype=torch.float32)
    yPredi = torch.tensor(ysDf[historyRecordNum:, ], dtype=torch.float32)

    xTrainDataloader = DataLoader(xTrain, batch_size=batchSize, shuffle=True, num_workers=8, drop_last=True)
    yTrainDataloader = DataLoader(yTrain, batch_size=batchSize, shuffle=True, num_workers=8, drop_last=True)
    xPrediDataloader = DataLoader(xPredi, batch_size=batchSize, shuffle=True, num_workers=8, drop_last=True)
    yPrediDataloader = DataLoader(yPredi, batch_size=batchSize, shuffle=True, num_workers=8, drop_last=True)

    return xTrainDataloader, yTrainDataloader, xPrediDataloader, yPrediDataloader
